import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np
import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'
#数据集：字符序列预测（Hello->Elloh)
char_set=list('hello')
char_to_idx={c:i for i,c in enumerate(char_set)}
idx_to_char={i:c for i,c in enumerate(char_set)}

#数据准备
input_str='hello'
target_str='elloh'
input_data=[char_to_idx[c] for c in input_str]
target_data=[char_to_idx[c] for c in target_str]
print(input_data)
print(target_data)
#转换为独热编码
input_one_hot=np.eye(len(char_set))[input_data]
print('one_hot:',input_one_hot)
#转换为PyTorch Tensor
inputs=torch.tensor(input_one_hot,dtype=torch.float32)
targets=torch.tensor(target_data,dtype=torch.long)

#模型超参数
input_size=len(char_set)
hidden_size=8
output_size=len(char_set)
num_epochs=200
learning_rate=0.1

#定义RNN模型
class RNNModel(nn.Module):
    def __init__(self,input_size,hidden_size,output_size):
        super(RNNModel,self).__init__()
        self.rnn=nn.RNN(input_size,hidden_size,batch_first=True)
        self.fc=nn.Linear(hidden_size,output_size)
    def forward(self,x,hidden):
        out,hidden=self.rnn(x,hidden)
        out=self.fc(out)
        return out,hidden

model=RNNModel(input_size,hidden_size,output_size)

#定义损失函数和优化器
criterion=nn.CrossEntropyLoss()
optimizer=optim.Adam(model.parameters(),lr=learning_rate)

#训练RNN
losses=[]
hidden=None
for epoch in range(num_epochs):
    optimizer.zero_grad()
    #前向传播
    outputs,hidden=model(inputs.unsqueeze(0),hidden)
    hidden=hidden.detach()#防止梯度爆炸

    #计算损失
    loss=criterion(outputs.view(-1,output_size),targets)
    loss.backward()
    optimizer.step()
    losses.append(loss.item())

    if (epoch+1)%20 ==0:
        print(f'Epoch[{epoch+1}/{num_epochs}],loss:{loss.item():.4f}')
with torch.no_grad():
    test_hidden=None
    test_output,_=model(inputs.unsqueeze(0),test_hidden)
    predicted=torch.argmax(test_output,dim=2).squeeze().numpy()
    print("Input sequence:",''.join([idx_to_char[i] for i in input_data]))
    print('Predicted sequence:',''.join([idx_to_char[i] for i in predicted]))

plt.plot(losses,label='Training Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title("RNN Training Loss Over Epochs")
plt.legend()
plt.show()
